import math
import numpy as np
import matplotlib.pyplot as plt 

# max_steps = 100000
# factor = 1.0
# model_size = 768

max_steps = 100000
warmup_steps = 4000
init_lr = 1.0e-7
peak_lr = 5.0e-4
final_lr = 1.0e-6
decay_start = 8000
decay_end = 80000


def get_linear_lr(i, start, end, start_lr, end_lr):
    assert end > start
    lr = start_lr + (i - start) * (end_lr - start_lr) / (end - start)
    return float(np.where(i < start, start_lr, np.where(i > end, end_lr, lr)))


def get_linear_warmup_peak_exp_decay_lr(step, lr):
    return float(np.where(
        step < warmup_steps,
        get_linear_lr(step, 0, warmup_steps, init_lr, peak_lr),
        np.where(
            step > decay_start,
            max(math.exp(get_linear_lr(step, decay_start, decay_end, math.log(peak_lr), math.log(final_lr))), final_lr),
            peak_lr
        )
    ))



# def get_lr(step):
#     return factor * model_size ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))

x = list(range(1, max_steps))

lrs = []
lr = init_lr
for i in range(1, max_steps):
    lr = get_linear_warmup_peak_exp_decay_lr(i, lr)
    lrs.append(lr)

plt.plot(x, lrs)
plt.savefig('lr.jpg')